import time
from unittest.mock import patch, MagicMock
import pytest

from motor.controller.core.instance_manager import InstanceManager
from motor.resources.endpoint import Endpoint, EndpointStatus
from motor.resources.http_msg_spec import HeartbeatMsg
from motor.resources.instance import ParallelConfig, Instance, NodeManagerInfo, InsStatus


class TestInstanceManager:

    def setup_method(self) -> None:
        """Setup for each test method"""
        # Clear singleton instance before each test
        if hasattr(InstanceManager, '_instances') and InstanceManager in InstanceManager._instances:
            try:
                InstanceManager._instances[InstanceManager].stop()
            except:
                pass
            if InstanceManager in InstanceManager._instances:
                del InstanceManager._instances[InstanceManager]
        
        self.dp = 8
        self.tp = 2
        self.p_role = "prefill"
        self.d_role = "decode"
        self.pod_ip1 = "127.0.0.1"
        self.pod_ip2 = "127.0.0.2"
        self.pod_ip3 = "127.0.0.3"
        self.pod_ip4 = "127.0.0.4"

        self.pod_ip5 = "127.0.0.5"
        self.pod_ip6 = "127.0.0.6"
        self.pod_ip7 = "127.0.0.7"
        self.pod_ip8 = "127.0.0.8"

        self.p_parallel_config = ParallelConfig(dp=self.dp, tp=self.tp)
        self.d_parallel_config = ParallelConfig(dp=self.dp * 4, tp=self.tp / 2)

    @pytest.fixture
    def instance_manager(self):
        """Setup mock instance manager"""
        with patch('threading.Thread') as mock_thread_class:
            mock_thread = MagicMock()
            mock_thread_class.return_value = mock_thread

            # add instance, 2P1D
            instance_manager = InstanceManager()
            # p0
            instance_manager.add_instance(Instance(
                job_name="prefill-0",
                model_name="test_model",
                id=0,
                role=self.p_role,
                parallel_config=self.p_parallel_config,
                node_mgrs=[NodeManagerInfo(pod_ip=self.pod_ip1, host_ip=self.pod_ip1, port="8080"),
                           NodeManagerInfo(pod_ip=self.pod_ip2, host_ip=self.pod_ip2, port="8080")],
                endpoints={self.pod_ip1: {
                    0: Endpoint(0, self.pod_ip1, port="9090", status=EndpointStatus.INITIAL, device_infos=list(),
                                hb_timestamp=time.time())},
                    self.pod_ip2: {
                        0: Endpoint(0, self.pod_ip2, port="9090", status=EndpointStatus.INITIAL, device_infos=list(),
                                    hb_timestamp=time.time())}}
            ))
            # p1
            instance_manager.add_instance(Instance(
                job_name="prefill-1",
                model_name="test_model",
                id=1,
                role=self.p_role,
                parallel_config=self.p_parallel_config,
                node_mgrs=[NodeManagerInfo(pod_ip=self.pod_ip3, host_ip=self.pod_ip3, port="8080"),
                           NodeManagerInfo(pod_ip=self.pod_ip4, host_ip=self.pod_ip4, port="8080")],
                endpoints={self.pod_ip3: {
                    0: Endpoint(0, self.pod_ip3, port="9090", status=EndpointStatus.INITIAL, device_infos=list(),
                                hb_timestamp=time.time())},
                    self.pod_ip4: {
                        0: Endpoint(0, self.pod_ip4, port="9090", status=EndpointStatus.INITIAL, device_infos=list(),
                                    hb_timestamp=time.time())}}
            ))
            # d0
            d_instance = Instance(
                job_name="decode-0",
                model_name="test_model",
                id=2,
                role=self.d_role,
                parallel_config=self.d_parallel_config,
                node_mgrs=[NodeManagerInfo(pod_ip=self.pod_ip5, host_ip=self.pod_ip5, port="8080"),
                           NodeManagerInfo(pod_ip=self.pod_ip6, host_ip=self.pod_ip6, port="8080"),
                           NodeManagerInfo(pod_ip=self.pod_ip7, host_ip=self.pod_ip7, port="8080"),
                           NodeManagerInfo(pod_ip=self.pod_ip8, host_ip=self.pod_ip8, port="8080"),
                           ],
                endpoints={}
            )
            # construct endpoints
            endpoints = {}
            for pod_ip in [self.pod_ip5, self.pod_ip6, self.pod_ip7, self.pod_ip8]:
                port_temp = 8080
                endpoints[pod_ip] = {}
                for i in range(0, 8):
                    endpoints[pod_ip][i] = Endpoint(
                        id=i,
                        ip=pod_ip,
                        port=str(port_temp),
                        status=EndpointStatus.INITIAL,
                        device_infos=[],
                        hb_timestamp=time.time()
                    )
                    port_temp += 1

                d_instance.add_endpoints(pod_ip, endpoints[pod_ip])

            instance_manager.add_instance(d_instance)
            return instance_manager

    def test_add_instance(self, instance_manager) -> None:
        """Test adding an instance"""
        cur_instance_num = instance_manager.get_instance_num()
        # Abnormal situation
        instance_manager.add_instance(self.p_parallel_config)
        actual_instance_num = instance_manager.get_instance_num()
        assert actual_instance_num == cur_instance_num

        instance_manager.add_instance(Instance(
            job_name="testAllocInsGroup2",
            model_name="test_model",
            id=100,
            role=self.p_role,
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp / 2)
        ))
        actual_instance_num = instance_manager.get_instance_num()
        assert actual_instance_num == cur_instance_num + 1

    def test_del_instance(self, instance_manager) -> None:
        """Test deleting an instance"""
        cur_instance_num = instance_manager.get_instance_num()

        instance_manager.del_instance(0)
        actual_instance_num = instance_manager.get_instance_num()
        assert actual_instance_num == cur_instance_num - 1

    def test_get_active_instances(self, instance_manager) -> None:
        """Test getting active instances"""
        instance = instance_manager.get_instance(0)
        instance.status = InsStatus.ACTIVE
        active_instances = instance_manager.get_active_instances()
        assert len(active_instances) == 1

    def test_get_inactive_instance(self, instance_manager) -> None:
        """Test getting inactive instances"""
        instance = instance_manager.get_instance(0)
        instance.status = InsStatus.INACTIVE
        inactive_instances = instance_manager.get_inactive_instances()
        assert len(inactive_instances) == 1

    def test_get_initial_instances(self, instance_manager) -> None:
        """Test getting initial instance status"""
        inactive_instances = instance_manager.get_initial_instances()
        assert len(inactive_instances) == 3

    def get_mock_heartbeat_msg_for_pinstance_normal(self, job_name: str, ins_id: int, ip: str) -> HeartbeatMsg:
        """Generate a mock heartbeat message"""
        status = {}
        for i in range(1):
            status[i] = EndpointStatus.NORMAL
        return HeartbeatMsg(
            job_name=job_name,
            ins_id=ins_id,
            ip=ip,
            status=status
        )

    def test_handle_heartbeat(self, instance_manager) -> None:
        """Test handling heartbeat"""
        # P0 ready
        mock_heartbeat_msg1 = self.get_mock_heartbeat_msg_for_pinstance_normal(
            "prefill-0",
            0,
            self.pod_ip1
        )
        instance_manager.handle_heartbeat(mock_heartbeat_msg1)
        instance = instance_manager.get_instance(0)
        assert instance.status == InsStatus.INITIAL

        mock_heartbeat_msg2 = self.get_mock_heartbeat_msg_for_pinstance_normal(
            "prefill-0",
            0,
            self.pod_ip2
        )
        instance_manager.handle_heartbeat(mock_heartbeat_msg2)
        instance = instance_manager.get_instance(0)
        assert instance.status == InsStatus.ACTIVE

        mock_heartbeat_msg3 = self.get_mock_heartbeat_msg_for_pinstance_normal(
            "prefill-0",
            0,
            self.pod_ip2
        )
        mock_heartbeat_msg3.status[0] = EndpointStatus.ABNORMAL
        instance_manager.handle_heartbeat(mock_heartbeat_msg3)
        instance = instance_manager.get_instance(0)
        assert instance.status == InsStatus.INACTIVE

        mock_heartbeat_msg4 = self.get_mock_heartbeat_msg_for_pinstance_normal(
            "prefill-1",
            1,
            self.pod_ip3
        )
        mock_heartbeat_msg4.status[0] = EndpointStatus.ABNORMAL
        instance_manager.handle_heartbeat(mock_heartbeat_msg4)
        instance = instance_manager.get_instance(1)
        assert instance is None

    def test_find_instance_with_matching_ip(self, instance_manager):
        """test finding an instance with matching ip"""
        pod_ip = "127.0.0.1"
        result = instance_manager.get_instance_by_podip(pod_ip)

        assert result is not None


    def test_no_instance_contains_ip(self, instance_manager):
        """test finding an instance with no matching ip"""
        pod_ip = "192.168.1.100"
        result = instance_manager.get_instance_by_podip(pod_ip)

        assert result is None

    def test_empty_string_pod_ip(self, instance_manager):
        """test empty string pod ip"""
        pod_ip = ""
        result = instance_manager.get_instance_by_podip(pod_ip)

        assert result is None